import numpy

from physics.vector.array import _array 
from physics.vector.array import Vector3Array

class Vector4Array(object):
    '''
    NumPy-like array of 4-vectors.
    
    Parameters
    ----------
    w : sequence of 4-tuples of length N
        A sequence of (x,y,z,t) sequences to be stored in the array

    OR 

    x : sequence of length N | scalar
        The x-components of the vectors to be stored in the array
    y : sequence of length N | scalar
        The y-components of the vectors to be stored in the array
    z : sequence of length N | scalar
        The z-components of the vectors to be stored in the array
    t : sequence of length N | scalar
        The t-components of the vectors to be stored in the array

    
    Optional Parameters
    -------------------
    dtype: NumPy datatype object
        The data type used for storage of the vector components (default is `numpy.float64')
    nan_to_num: bool
        if True, all functions will return arrays filtered through NumPy `nan_to_num' function to 
        remove nans and infs.

    '''

    def __init__(self, *args, **kwargs):
        # TODO:
        # - add support for setting r, theta, phi
        # - add support for setting pt, eta, phi
        # - add support for setting pt, theta, phi

        if 'dtype' in kwargs:
            self.dtype = kwargs['dtype']
        else:
            self.dtype = numpy.float64

        if 'nan_to_num' in kwargs:
            self.nan_to_num = kwargs['nan_to_num']
        else:
            self.nan_to_num = False

        if len(args) == 1:
            if sum(map(len, args[0])) != 4*len(args[0]):
                msg = "Calling Vector4Array.__init__() with a single argument requires that argument to be"
                msg += " a sequence of 4-tuples"
                raise ValueError(msg)
    
            x, y, z, t = zip(*args[0])

        elif len(args) == 3:
            x, y, z = args

            if 'm' in kwargs:
                t = numpy.sqrt(x**2 + y**2 + z**2 + kwargs['m']**2)
            elif 'mass' in kwargs:
                t = numpy.sqrt(x**2 + y**2 + z**2 + kwargs['mass']**2)
            else:
                msg = "Calling Vector4Array.__init__() with three arguments requires the `m' or `mass' kwarg"
                raise ValueError(msg)
            
        elif len(args) == 4:
            x, y, z, t = args
        
        else:
            msg = "Calling Vector4Array.__init__() with {0} arguments is undefined".format(len(args))
            raise ValueError(msg)

        if 'length' in kwargs:
            self.v = Vector3Array(x, y, z, length=kwargs['length'], dtype=self.dtype)
        else:
            self.v = Vector3Array(x, y, z, dtype=self.dtype)
        
        self.x = self.v.x
        self.y = self.v.y
        self.z = self.v.z

        if numpy.isscalar(t):
            self.t = numpy.ones(self.v.shape, dtype=self.dtype) * t
        else:
            self.t = numpy.array(t,  dtype=self.dtype)

        if len(self.x.shape) > 1:
            msg = "Argument 1 has shape {0}; must be a one-dimensional sequence".format(self.x.shape)
            raise ValueError(msg)

        if len(self.y.shape) > 1:
            msg = "Argument 2 has shape {0}; must be a one-dimensional sequence".format(self.y.shape)
            raise ValueError(msg)

        if len(self.z.shape) > 1:
            msg = "Argument 3 has shape {0}; must be a one-dimensional sequence".format(self.z.shape)
            raise ValueError(msg)

        if len(self.t.shape) > 1:
            msg = "Argument 4 has shape {0}; must be a one-dimensional sequence".format(self.t.shape)
            raise ValueError(msg)

        self._check_component_length()

        self.shape = self.x.shape

    def phi(self):
        '''
        Returns an array containing the azimuthal angle on (-pi, pi) of the space 
        component of each 4-vector.
        '''
        
        if self.nan_to_num:
            return numpy.nan_to_num(numpy.arctan2(self.y, self.x))

        return numpy.arctan2(self.y, self.x)
        
    def theta(self):
        '''
        Returns an array containing the polar angle on (0, pi) of the space 
        component of each 4-vector.
        '''

        if self.nan_to_num:
            return numpy.nan_to_num(numpy.arctan2(self.pt(), self.z))

        return numpy.arctan2(self.pt(), self.z)
        
    def cos_theta(self):
        '''
        Returns an array containing the cosine of the polar angle of the space 
        component of each 4-vector.
        '''

        if self.nan_to_num:
            return numpy.nan_to_num(_array.v3_cos_theta(self.x, self.y, self.z))
        
        return _array.v3_cos_theta(self.x, self.y, self.z)


    def r2(self):
        '''
        Returns an array containing the magnitude squared of the space component of each 4-vector.
        '''

        return self.x**2 + self.y**2 + self.z**2

    def r(self):
        '''
        Returns an array containing the magnitude of the space component of each 4-vector.
        '''

        return numpy.sqrt(self.r2())


    def boost_vector(self):
        if numpy.count_nonzero(self.t) != len(self):
            warnings.warn("time component of Vector4Array contains zeros--boost vector will contain NaNs")
                   
        return Vector3Array(self.x/self.t, self.y/self.t, self.z/self.t)

    def boost(self, *args):
        if len(args) == 1:
            bv = args[0]
            
            if hasattr(bv, 'shape'):
                if bv.shape[0] != len(self):
                    if bv.shape[0] != 3:
                        msg = "Single argument supplied has shape `{0}'; expected (3,) or ({1},)"
                        raise ValueError(msg.format(bv.shape, len(self)))
                
                    _array.v4_boost_onevec(self.x, self.y, self.z, self.t, bv[0], bv[1], bv[2])

                else:            
                    _array.v4_boost_arr(self.x, self.y, self.z, self.t, bv.x, bv.y, bv.z)

            elif hasattr(bv, '__len__'):
                if len(bv) == 3:
                    _array.v4_boost_onevec(self.x, self.y, self.z, self.t, bv[0], bv[1], bv[2])
                    return                    
                else:
                    msg = "Single argument has improper length of {0}; should be 3 or {1}"
                    raise ValueError(msg.format(len(bv), len(self)))
            else:
                if hasattr(bv, 'X') and hasattr(bv, 'Y') and hasattr(bv, 'Z'):
                     _array.v4_boost_onevec(self.x, self.y, self.z, self.t, bv.X(), bv.Y(), bv.Z())
                else:
                    raise ValueError("Cannot boost with type `{0}'".format(type(args[0])))

        elif len(args) == 3:
            _array.v4_boost_onevec(self.x, self.y, self.z, self.t, args[0], args[1], args[2])
        else:
            msg = "Vector4Array.boost() called with unsupported number of arguments ({0})"
            raise ValueError(msg.format(len(args)))                

    def m2(self):
        return self.t**2 - self.x**2 - self.y**2 - self.z**2

    def m(self):
        return _array.unsafe_sqrt(self.m2())
    
    def pt(self):
        return self.v.pt()

    def pt2(self):
        return self.v.pt2()

    def et2(self, axis=None):
        if axis is None:
            return _array.v4_et2(self.x, self.y, self.z, self.t)            
        elif type(axis) == type(self.v):            
            if axis.shape != self.v.shape:
                raise ValueError("Axis shape {0} doesn't match data shape {1}".format(axis.shape, data.shape))
            return _array.v4_et2_multi_axis(self.x, self.y, self.z, self.t, 
                                            axis.x, axis.y, axis.z)
        elif len(axis) == 3:
            return _array.v4_et2_single_axis(self.x, self.y, self.z, self.t,
                                             float(axis[0]), float(axis[1]), float(axis[2]))
        else:
            msg = "`axis' argument must be a Vector3Array of shape `{0}', a 3-tuple, or None"
            raise ValueError(msg.format(self.shape))
        
    def et(self, axis=None):
        return numpy.sqrt(self.et2(axis))

    def beta(self):
        return self.v.mag() / self.t

    def gamma(self):
        return _array.v4_gamma(self.x, self.y, self.z, self.t)
    
    def __add__(self, rhs):
        if type(rhs) != type(self):
            raise ValueError("Cannot add type `{0}' to type `{1}'".format(type(rhs),
                                                                          type(self)))

        return Vector4Array(self.x + rhs.x, self.y + rhs.y, 
                            self.z + rhs.z, self.t + rhs.t)

    def __iadd__(self, rhs):
        if type(rhs) != type(self):
            raise ValueError("Cannot add type `{0}' to type `{1}'".format(type(rhs),
                                                                          type(self)))

        self.x += rhs.x
        self.y += rhs.y
        self.z += rhs.z
        self.t += rhs.t

        return self
        

    def __sub__(self, rhs):
        if type(rhs) != type(self):
            raise ValueError("Cannot subtract type `{0}' from type `{1}'".format(type(rhs),
                                                                                 type(self)))

        return Vector4Array(self.x - rhs.x, self.y - rhs.y, 
                            self.z - rhs.z, self.t - rhs.t)
        
    
    def __isub__(self, rhs):
        if type(rhs) != type(self):
            raise ValueError("Cannot subtract type `{0}' from type `{1}'".format(type(rhs),
                                                                                 type(self)))

        self.x -= rhs.x
        self.y -= rhs.y
        self.z -= rhs.z
        self.t -= rhs.t

        return self    

    def __mul__(self, rhs):
        if numpy.isscalar(rhs):
            return Vector4Array(rhs*self.x,
                                rhs*self.y,
                                rhs*self.z,
                                rhs*self.t)

        elif type(rhs) == type(self):
            return self.dot(rhs)
    
        else:
            msg = "Multiplication undefined between types `{0}' and `{1}'".format(type(rhs),
                                                                                  type(self))
            raise ValueError(msg)

    def __imul__(self, rhs):
        if not numpy.isscalar(rhs):
            msg = "In-place multiplication undefined between types `{0}' and `{1}'"
            raise ValueError(msg.format(type(rhs), type(self)))
            
        self.x *= rhs
        self.y *= rhs
        self.z *= rhs
        self.t *= rhs

        return self    

    def __div__(self, rhs):
        if not numpy.isscalar(rhs):
            msg = "Division undefined between types `{0}' and `{1}'"
            raise ValueError(msg.format(type(rhs), type(self)))

        return Vector4Array(self.x/rhs,
                            self.y/rhs,
                            self.z/rhs,
                            self.t/rhs)
    

    def __idiv__(self, rhs):
        if not numpy.isscalar(rhs):
            msg = "In-place division undefined between types `{0}' and `{1}'"
            raise ValueError(msg.format(type(rhs), type(self)))
            
        self.x /= rhs
        self.y /= rhs
        self.z /= rhs
        self.t /= rhs

        return self    

    
    def __getitem__(self, index):
        return self.x[index], self.y[index], self.z[index]

    def __len__(self):
        return self.shape[0]
            
    def __eq__(self, rhs):
        if type(rhs) != type(self):
            raise ValueError("Cannot compare type `{0}' to type `{1}'".format(type(rhs),
                                                                              type(self)))
        
        return self.x == rhs.x and self.y == rhs.y and self.z == rhs.z and self.t == rhs.t

    def __ne__(self, rhs):
        if type(rhs) != type(self):
            raise ValueError("Cannot compare type `{0}' to type `{1}'".format(type(rhs),
                                                                              type(self)))

        return self.x != rhs.x and self.y != rhs.y and self.z != rhs.z and self != rhs.t

    def __neg__(self):
        return Vector4Array(-self.x, -self.y, -self.z, -self.t)

    def __str__(self):
        print "POOPIE"
        return str(zip(self.x, self.y, self.z, self.e))

    def __iter__(self):
        for i in xrange(self.shape[0]):
            yield (self.x[i], self.y[i], self.z[i])            
    
    def _check_component_length(self):
        import operator
        import itertools

        lengths = dict(zip('xyzt', map(len, [self.x, self.y, self.z, self.t])))
        length_comp = [operator.eq(*pair) for pair in itertools.combinations(lengths.values(), 2)]

        if not numpy.all(length_comp):
            msg = "Arguments do not all have the same length: {0}".format(lengths)
            raise ValueError(msg)

    
